#include "VolumeRenderer.h"

bool ColorPointLT(const ColorPoint& p1, const ColorPoint& p2)
{
    return p1.x < p2.x;
}

bool TFEditor::save()
{
	std::ofstream fout("TF.txt");
	if(!fout.is_open())
		return false;
	fout << points.size() << std::endl;
	for(int i = 0; i < points.size(); ++i)
		fout << points[i].x << ' ' << int(points[i].y.r) << ' ' << int(points[i].y.g) << ' ' << int(points[i].y.b) << std::endl;
	fout.close();
	function.save("TF.png");
	return true;
}

bool TFEditor::load()
{
	std::ifstream fin("TF.txt");
	if(!fin.is_open())
		return false;
	int np, r, g, b;
	fin >> np;
	points.resize(np);
	for(int i = 0; i < points.size(); ++i)
	{
		points[i].sel = false;
		fin >> points[i].x >> r >> g >> b;
		points[i].y = Color3uint8(r, g, b);
	}
	fin.close();
	function.load("TF.png");
	update();
	return true;
}

void TFEditor::insert(float x, const Color3uint8 &y)
{
	assert(x >= 0.0f && x <= 1.0f);
	points.append(ColorPoint(x, y));
}

void TFEditor::update()
{
	if(points.size() == 0)
	{
		for(int i = 0; i < function.width(); ++i)
			function.pixel4(i, 0) = Color4uint8(0, 0, 0, function.pixel4(i, 0).a);
		return;
	}

	for(int i = 0; i < points.size(); ++i)
		if(points[i].sel)
			points[i].y = pick;

	points.sort(ColorPointLT);

	int k = 0;
	for(int i = 0; i < function.width(); ++i)
	{
		float x = (i + 0.5f) / function.width();
		if(x < points.front().x)
		{
			function.pixel4(i, 0) = Color4uint8(points.front().y, function.pixel4(i, 0).a);
			continue;
		}
		else if(x > points.back().x)
		{
			function.pixel4(i, 0) = Color4uint8(points.back().y, function.pixel4(i, 0).a);
			continue;
		}
		else
		{
			while(x > points[k + 1].x)
				++k;
			function.pixel4(i, 0) = Color4uint8(lerp(points[k].y, points[k + 1].y, (x - points[k].x) / (points[k + 1].x - points[k].x)), function.pixel4(i, 0).a);
		}
	}

	colormap = Texture::fromGImage("", function, ImageFormat::RGBA8(), Texture::DIM_2D_NPOT, Texture::Settings::video());
}

template<typename T>
void TFEditor::stat(const Array<T> &data, T minVal, T maxVal, float scale)
{
	float p = scale / data.size();
	for(int i = 0; i < histogram.size(); ++i)
		histogram[i] = 0.0f;
	for(int i = 0; i < data.size(); ++i)
	{
		int index = iClamp(int(histogram.size() * (data[i] - minVal) / (maxVal - minVal)), 0, histogram.size() - 1);
		histogram[index] += p;
	}
	for(int i = 0; i < histogram.size(); ++i)
		histogram[i] = clamp(log(histogram[i]) * scale + 1.0f, 0.0f, 1.0f);
}

void TFEditor::render(RenderDevice *rd, const GuiThemeRef &skin) const
{
	skin->pauseRendering();
	Draw::fastRect2D(m_rect, rd, Color3::black());

	// histogram
	Rect2D rect1 = Rect2D::xywh(m_rect.x0(), m_rect.y0(), m_rect.width(), m_rect.height() * layoutT);
	//Draw::fastRect2D(rect1, rd, Color3::blue());
	rd->setColor(Color3::gray());
	rd->setLineWidth(1.0f);
	rd->beginPrimitive(RenderDevice::Primitive::LINES);
	for(int i = 1; i < 10; ++i)
	{
		int y = rect1.y0() + rect1.height() * i / 10;
		rd->sendVertex(Vector2(rect1.x0(), y));
		rd->sendVertex(Vector2(rect1.x1(), y));
	}
	for(int i = 1; i < 10; ++i)
	{
		int x = rect1.x0() + rect1.width() * i / 10;
		rd->sendVertex(Vector2(x, rect1.y0()));
		rd->sendVertex(Vector2(x, rect1.y1()));
	}
	rd->endPrimitive();
	rd->setColor(Color3::yellow());
	rd->setLineWidth(2.0f);
	rd->beginPrimitive(RenderDevice::Primitive::LINES);
	for(int i = 0; i < histogram.size(); ++i)
	{
		int x = rect1.x0() + (rect1.width() - 1) * i / (histogram.size() - 1);
		rd->sendVertex(Vector2(x, rect1.y1() - histogram[i] * rect1.height()));
		rd->sendVertex(Vector2(x, rect1.y1()));
	}
	rd->endPrimitive();
	rd->setColor(Color3::green());
	rd->setLineWidth(2.0f);
	rd->beginPrimitive(RenderDevice::Primitive::LINE_STRIP);
	for(int x = rect1.x0(); x <= rect1.x1(); ++x)
	{
		int v = iClamp(int(function.width() * (x - rect1.x0()) / rect1.width()), 0, function.width() - 1);
		int y = int(rect1.y1() + (rect1.y0() - rect1.y1()) * (function.pixel4(v, 0).a / 255.0f));
		rd->sendVertex(Vector2(x, y));
	}
	rd->endPrimitive();

	// colormap
	Rect2D rect2 = Rect2D::xyxy(m_rect.x0(), m_rect.y0() + m_rect.height() * (1.0f + layoutT - layoutB) * 0.5f, m_rect.x1(), m_rect.y1() - m_rect.height() * layoutB);
	//Draw::fastRect2D(rect2, rd, Color3::white());
	rd->setTexture(0, colormap);
	Draw::fastRect2D(rect2, rd);
	rd->setTexture(0, NULL);

	// spacemap
	Rect2D rect3 = Rect2D::xyxy(m_rect.x0(), m_rect.y1() - m_rect.height() * layoutB, m_rect.x1(), m_rect.y1());
	rd->setTexture(0, spacemap);
	Draw::fastRect2D(rect3, rd);
	rd->setTexture(0, NULL);

	// frames
	rd->setColor(Color3::white());
	rd->setLineWidth(1.0f);
	rd->beginPrimitive(RenderDevice::Primitive::LINES);
	rd->sendVertex(Vector2(m_rect.x0(), m_rect.y0()));
	rd->sendVertex(Vector2(m_rect.x0(), m_rect.y1()));
	rd->sendVertex(Vector2(m_rect.x0(), m_rect.y1()));
	rd->sendVertex(Vector2(m_rect.x1(), m_rect.y1()));
	rd->sendVertex(Vector2(m_rect.x1(), m_rect.y1()));
	rd->sendVertex(Vector2(m_rect.x1(), m_rect.y0()));
	rd->sendVertex(Vector2(m_rect.x1(), m_rect.y0()));
	rd->sendVertex(Vector2(m_rect.x0(), m_rect.y0()));
	rd->sendVertex(Vector2(rect1.x0(), rect1.y1()));
	rd->sendVertex(Vector2(rect1.x1(), rect1.y1()));
	rd->sendVertex(Vector2(rect2.x0(), rect2.y0()));
	rd->sendVertex(Vector2(rect2.x1(), rect2.y0()));
	rd->sendVertex(Vector2(rect2.x0(), rect2.y1()));
	rd->sendVertex(Vector2(rect2.x1(), rect2.y1()));
	rd->endPrimitive();

	// control points
	for(int i = 0; i < points.size(); ++i)
	{
		if(points[i].sel)
			continue;
		float w = rect2.height() * 0.5f;
		float h = w * sqrt(3.0f);
		Array<Vector2> tri;
		tri.resize(3);
		tri[0] = Vector2(rect2.x0() + points[i].x * rect2.width(), rect2.y0());
		tri[1] = tri[0] + Vector2(-w, -h) * 0.75f;
		tri[2] = tri[0] + Vector2(w, -h) * 0.75f;
		Draw::poly2D(tri, rd, Color4uint8(points[i].y, 255));
		Draw::poly2DOutline(tri, rd, Color3::white());
	}
	for(int i = 0; i < points.size(); ++i)
	{
		if(!points[i].sel)
			continue;
		float w = rect2.height() * 0.5f;
		float h = w * sqrt(3.0f);
		Array<Vector2> tri;
		tri.resize(3);
		tri[0] = Vector2(rect2.x0() + points[i].x * rect2.width(), rect2.y0());
		tri[1] = tri[0] + Vector2(-w, -h);
		tri[2] = tri[0] + Vector2(w, -h);
		Draw::poly2D(tri, rd, Color4uint8(points[i].y, 255));
		Draw::poly2DOutline(tri, rd, Color3::white());
	}

	skin->resumeRendering();
}

bool TFEditor::onEvent(const GEvent &event)
{
	Rect2D rect1 = Rect2D::xywh(m_rect.x0(), m_rect.y0(), m_rect.width(), m_rect.height() * layoutT);
	Rect2D rect2 = Rect2D::xyxy(m_rect.x0(), m_rect.y0() + m_rect.height() * (1.0f + layoutT - layoutB) * 0.5f, m_rect.x1(), m_rect.y1() - m_rect.height() * layoutB);
	Rect2D rect3 = Rect2D::xyxy(m_rect.x0(), m_rect.y1() - m_rect.height() * layoutB, m_rect.x1(), m_rect.y1());

	if((event.type == GEventType::KEY_DOWN) && (event.key.keysym.sym == GKey::DELETE))
	{
		for(int i = 0; i < points.size(); ++i)
		{
			if(points[i].sel)
				points.remove(i--);
		}

		update();
		return true;
	}

	if((event.type == GEventType::MOUSE_BUTTON_DOWN) && m_rect.contains(Vector2(event.button.x, event.button.y)))
	{
		if(rect1.contains(Vector2(event.button.x, event.button.y)))
		{
			mode = 1;	// draw
			x0 = iClamp(int16(event.button.x), rect1.x0(), rect1.x1());
			y0 = iClamp(int16(event.button.y), rect1.y0(), rect1.y1());
			return true;
		}
		else if(rect3.contains(Vector2(event.button.x, event.button.y)))
		{
			mode = 3;
			x0 = iClamp(int16(event.button.x), rect3.x0(), rect3.x1());
			y0 = iClamp(int16(event.button.y), rect3.y0(), rect3.y1());
			int x = iClamp(int(hslspace.width() * x0 / rect3.width()), 0, hslspace.width() - 1);
			int y = iClamp(int(hslspace.height() * (y0 - rect3.y0()) / rect3.height()), 0, hslspace.height() - 1);
			pick = hslspace.pixel3(x, y);
			for(int i = 0; i < points.size(); ++i)
			{
				if(points[i].sel)
					points[i].y = pick;
			}
			update();
			return true;
		}
		else
		{
			for(int i = 0; i < points.size(); ++i)
				points[i].sel = false;

			for(int i = 0; i < points.size(); ++i)
			{
				int r = int(rect2.height() * 0.5f);
				Vector2int16 point = Vector2int16(rect2.x0() + points[i].x * rect2.width(), rect2.y0() - r);
				if(abs(point.x - int16(event.button.x)) < r && abs(point.y - int16(event.button.y)) < r)
				{
					mode = 2;	// drag
					points[i].sel = true;
					pick = points[i].y;
					x0 = iClamp(int16(event.button.x), rect2.x0(), rect2.x1());
					y0 = iClamp(int16(event.button.y), rect2.y0(), rect2.y1());
					return true;
				}
			}

			x0 = iClamp(int16(event.button.x), rect2.x0(), rect2.x1());
			y0 = iClamp(int16(event.button.y), rect2.y0(), rect2.y1());
			mode = 2;
			insert((x0 - rect2.x0()) / rect2.width(), pick);
			points.last().sel = true;
			update();
			return true;
		}
	}

	if(event.type == GEventType::MOUSE_BUTTON_UP)
	{
		mode = 0;
		return true;
	}

	if(mode == 1 && (event.type == GEventType::MOUSE_MOTION))
	{
		int mx = iClamp(int16(event.button.x), rect1.x0(), rect1.x1());
		int my = iClamp(int16(event.button.y), rect1.y0(), rect1.y1());
		int x1 = mx;
		int y1 = my;

		if(x0 > x1)
		{
			int xt = x0;
			x0 = x1;
			x1 = xt;
			int yt = y0;
			y0 = y1;
			y1 = yt;
		}

		for(int x = x0; x <= x1; ++x)
		{
			int y = (x0 == x1 ? y1 : y0 + (y1 - y0) * (x - x0) / (x1 - x0));
			int v = iClamp(int(function.width() * (x - rect1.x0()) / rect1.width()), 0, function.width() - 1);
			int c = iClamp(int(256 * (int(rect1.y1()) - y) / rect1.height()), 0, 255);
			function.pixel4(v, 0).a = c;
		}

		x0 = mx;
		y0 = my;
		colormap = Texture::fromGImage("", function, ImageFormat::RGBA8(), Texture::DIM_2D_NPOT, Texture::Settings::video());
		save();
		return true;
	}

	if(mode == 2 && (event.type == GEventType::MOUSE_MOTION))
	{
		int mx = iClamp(int16(event.button.x), rect2.x0(), rect2.x1());
		int my = iClamp(int16(event.button.y), rect2.y0(), rect2.y1());
		for(int i = 0; i < points.size(); ++i)
		{
			if(points[i].sel)
			{
				points[i].x += float(mx - x0) / rect2.width();
				points[i].x = clamp(points[i].x, 0.0f, 1.0f);
			}
		}
		x0 = mx;
		y0 = my;
		update();
		return true;
	}

	if(mode == 3 && (event.type == GEventType::MOUSE_MOTION))
	{
		x0 = iClamp(int16(event.button.x), rect3.x0(), rect3.x1());
		y0 = iClamp(int16(event.button.y), rect3.y0(), rect3.y1());
		int x = iClamp(int(hslspace.width() * x0 / rect3.width()), 0, hslspace.width() - 1);
		int y = iClamp(int(hslspace.height() * (y0 - rect3.y0()) / rect3.height()), 0, hslspace.height() - 1);
		pick = hslspace.pixel3(x, y);
		for(int i = 0; i < points.size(); ++i)
		{
			if(points[i].sel)
				points[i].y = pick;
		}
		update();
		return true;
	}

	return false;
}

TFEditor::TFEditor(GuiContainer *parent, const GuiText &text, int width, int height) : GuiControl(parent, text)
{
	layoutT = 0.6f;
	layoutB = 0.3f;
	mode = 0;
	x0 = y0 = -1;
	setWidth(width);
	setHeight(height);
	function.resize(256, 1, 4, true);
	histogram.resize(function.width());
	memset(&histogram[0], 0, histogram.size() * sizeof(float));
	hslspace.resize(384, 128, 3, true);
	int period = hslspace.width() / 6;
	Color3 black(0.0f, 0.0f, 0.0f);
	Color3 white(1.0f, 1.0f, 1.0f);
	Color3 color[7];
	color[0] = Color3(1.0f, 0.0f, 0.0f);
	color[1] = Color3(1.0f, 1.0f, 0.0f);
	color[2] = Color3(0.0f, 1.0f, 0.0f);
	color[3] = Color3(0.0f, 1.0f, 1.0f);
	color[4] = Color3(0.0f, 0.0f, 1.0f);
	color[5] = Color3(1.0f, 0.0f, 1.0f);
	color[6] = Color3(1.0f, 0.0f, 0.0f);
	for(int i = 0; i < hslspace.width(); ++i)
	{
		int cyc = i / period;
		Color3 mid = lerp(color[cyc], color[cyc + 1], float(i - cyc * period) / period);
		int hh = hslspace.height() / 2;
		for(int j = 0; j < hh; ++j)
			hslspace.pixel3(i, j) = lerp(black, mid, float(j) / hh);
		for(int j = hh; j < hslspace.height(); ++j)
			hslspace.pixel3(i, j) = lerp(mid, white, float(j - hh) / hh);
	}
	spacemap = Texture::fromGImage("", hslspace);
	if(!load())
	{
		insert(0.0f, Color3uint8(0, 0, 255));
		insert(0.25f, Color3uint8(0, 255, 255));
		insert(0.5f, Color3uint8(0, 255, 0));
		insert(0.75f, Color3uint8(255, 255, 0));
		insert(1.0f, Color3uint8(255, 0, 0));
		update();
	}
}
